import json
from my_utils.api_request import *
from my_utils.data_type import QueryResult, QueryResultItem
from dacite import from_dict
def get_chunk_by_query(query,top_k=10):
    
    search_results = search_papers(query, top_k)
    wraped_search_result = QueryResult(results=[
    from_dict(data_class=QueryResultItem, data=item) 
    for item in search_results
    ])
    return_text = ""
    for i,item_result in enumerate(wraped_search_result.results):
        return_text += f"# 第{i+1}个切块\n## 切块编号（chunk_id）:{item_result.entity.chunk_id}\n## 论文编号（paper_id）:{item_result.entity.paper_id}\n## 论文标题:{item_result.entity.paper_title}\n## 论文切块内容:\n{item_result.entity.chunk_text}\n\n"
    # print(return_text)
    return return_text

    
    

# 根据sql语句查询数据库
tool_get_chunk_by_query = {
    "tool_schema":{
        "type": "function",
        "function": {
            "name": "get_chunk_by_query",
            "strict": True,
            "description": "基于文献查询语句，执行模糊查询，最终返回相关的论文切块",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "英文查询语句",
                    },
                    "top_k": {
                        "type": "number",
                        "description": "返回的论文切块数量",
                    }
                },
                "required": ["query", "top_k"],
                "additionalProperties": False,
            },
        },
    },
    "tool_def": get_chunk_by_query
}

if __name__ == "__main__":
    sql_result = get_chunk_by_query("机器学习")
    print(sql_result)